import numpy as np
from time import time
import pyopencl as cl
import pyopencl.array
from pyopencl.elementwise import ElementwiseKernel


def run(n):
    print(__file__)
    print("n:", n)

    rng = np.random.default_rng()
    a_np = rng.random(n, dtype=np.float32)
    b_np = rng.random(n, dtype=np.float32)

    ctx = cl.create_some_context()
    queue = cl.CommandQueue(ctx)

    a_g = cl.array.to_device(queue, a_np)
    b_g = cl.array.to_device(queue, b_np)

    lin_comb = ElementwiseKernel(ctx,
                                 "float k1, float *a_g, float k2, float *b_g, float *res_g",
                                 "res_g[i] = k1 * a_g[i] + k2 * b_g[i]",
                                 "lin_comb")

    res_g = cl.array.empty_like(a_g)
    tcl1 = time()
    lin_comb(2, a_g, 3, b_g, res_g)
    tcl2 = time()

    # Check on GPU with PyOpenCL Array:
    (res_g - (2 * a_g + 3 * b_g)).get()

    # Check on CPU with Numpy:
    tnp1 = time()
    np_res = 2 * a_np + 3 * b_np
    tnp2 = time()
    res_np = res_g.get()
    res_np - (2 * a_np + 3 * b_np)
    print(np.linalg.norm(res_np - (2 * a_np + 3 * b_np)))

    # Check on CPU with native python:
    tna1 = time()
    na_res = [2 * a_np[i] + 3 * b_np[i] for i in range(n)]
    tna2 = time()

    # time
    print(f"Python : {tna2 - tna1:.10f}s")
    print(f"Numpy  : {tnp2 - tnp1:.10f}s")
    print(f"OpenCL : {tcl2 - tcl1:.10f}s")


if __name__ == '__main__':
    run(10)
    run(100)
    run(1000)
    run(10000)
    run(100000)
    run(1000000)
    run(10000000)
